package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/google/uuid"
	"github.com/shirou/gopsutil/v3/cpu"
	"github.com/shirou/gopsutil/v3/disk"
	"github.com/shirou/gopsutil/v3/host"
	"github.com/shirou/gopsutil/v3/mem"
)

// AgentConfig 代理配置
type AgentConfig struct {
	MasterURL    string
	Token        string
	NodeName     string
	NodeIP       string
	PollInterval time.Duration
	Insecure     bool
}

// WorkerRegistration Worker注册信息
type WorkerRegistration struct {
	ID        string `json:"id,omitempty"`
	Name      string `json:"name"`
	IPAddress string `json:"ip_address"`
	Token     string `json:"token,omitempty"`
}

// HeartbeatRequest 心跳请求
type HeartbeatRequest struct {
	Token string `json:"token"`
}

// TaskInfo 任务信息
type TaskInfo struct {
	ID        string      `json:"id"`
	TestCaseID string     `json:"test_case_id"`
	FIOConfig interface{} `json:"fio_config"`
}

// TaskResult 任务结果
type TaskResult struct {
	Status  string      `json:"status"`
	JobName string      `json:"jobname,omitempty"`
	Read    interface{} `json:"read,omitempty"`
	Write   interface{} `json:"write,omitempty"`
	Latency interface{} `json:"latency,omitempty"`
	Error   string      `json:"error,omitempty"`
}

// SystemInfo 系统信息
type SystemInfo struct {
	OS        string  `json:"os"`
	Platform  string  `json:"platform"`
	Arch      string  `json:"arch"`
	CPUs      int     `json:"cpus"`
	Memory    uint64  `json:"memory_mb"`
	DiskSpace uint64  `json:"disk_space_gb"`
	CPUUsage  float64 `json:"cpu_usage_percent"`
	MemUsage  float64 `json:"mem_usage_percent"`
}

// Agent xfio代理
type Agent struct {
	config       *AgentConfig
	client       *http.Client
	workerID     string
	workerToken  string
	registered   bool
	mu           sync.RWMutex
	taskChan     chan TaskInfo
	taskCancel   context.CancelFunc
	taskContext  context.Context
}

// NewAgent 创建新代理
func NewAgent(config *AgentConfig) *Agent {
	// 创建HTTP客户端
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: config.Insecure},
	}
	
	client := &http.Client{
		Transport: tr,
		Timeout:   30 * time.Second,
	}
	
	// 获取本机IP地址（如果没有指定）
	nodeIP := config.NodeIP
	if nodeIP == "" {
		nodeIP = getLocalIP()
	}
	
	// 获取节点名称（如果没有指定）
	nodeName := config.NodeName
	if nodeName == "" {
		nodeName = fmt.Sprintf("xfio-node-%s", uuid.New().String()[:8])
	}
	
	config.NodeIP = nodeIP
	config.NodeName = nodeName
	
	return &Agent{
		config:   config,
		client:   client,
		taskChan: make(chan TaskInfo, 10),
	}
}

// getLocalIP 获取本机IP地址
func getLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return "127.0.0.1"
	}
	
	for _, addr := range addrs {
		if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	
	return "127.0.0.1"
}

// RegisterWithMaster 向主节点注册
func (a *Agent) RegisterWithMaster() error {
	registration := WorkerRegistration{
		Name:      a.config.NodeName,
		IPAddress: a.config.NodeIP,
	}
	
	jsonData, err := json.Marshal(registration)
	if err != nil {
		return fmt.Errorf("failed to marshal registration data: %v", err)
	}
	
	url := fmt.Sprintf("%s/api/workers/register", strings.TrimRight(a.config.MasterURL, "/"))
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create registration request: %v", err)
	}
	
	req.Header.Set("Content-Type", "application/json")
	
	resp, err := a.client.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send registration request: %v", err)
	}
	defer resp.Body.Close()
	
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("registration failed with status %d: %s", resp.StatusCode, string(body))
	}
	
	var result WorkerRegistration
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return fmt.Errorf("failed to decode registration response: %v", err)
	}
	
	a.mu.Lock()
	a.workerID = result.ID
	a.registered = true
	a.mu.Unlock()
	
	log.Printf("Successfully registered with master. Worker ID: %s", a.workerID)
	return nil
}

// SendHeartbeat 发送心跳
func (a *Agent) SendHeartbeat() error {
	a.mu.RLock()
	if !a.registered {
		a.mu.RUnlock()
		return fmt.Errorf("agent not registered")
	}
	
	a.mu.RUnlock()
	
	url := fmt.Sprintf("%s/api/workers/%s/heartbeat", strings.TrimRight(a.config.MasterURL, "/"), a.workerID)
	
	// 创建心跳请求
	heartbeatReq := HeartbeatRequest{
		Token: a.workerToken,
	}
	
	jsonData, err := json.Marshal(heartbeatReq)
	if err != nil {
		return fmt.Errorf("failed to marshal heartbeat data: %v", err)
	}
	
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create heartbeat request: %v", err)
	}
	
	req.Header.Set("Content-Type", "application/json")
	
	resp, err := a.client.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send heartbeat: %v", err)
	}
	defer resp.Body.Close()
	
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("heartbeat failed with status %d: %s", resp.StatusCode, string(body))
	}
	
	log.Printf("Heartbeat sent successfully")
	return nil
}

// GetSystemInfo 获取系统信息
func (a *Agent) GetSystemInfo() (*SystemInfo, error) {
	info := &SystemInfo{
		OS:       runtime.GOOS,
		Arch:     runtime.GOARCH,
		CPUs:     runtime.NumCPU(),
	}
	
	// 获取平台信息
	hostInfo, err := host.Info()
	if err == nil {
		info.Platform = hostInfo.Platform
	}
	
	// 获取内存信息
	memInfo, err := mem.VirtualMemory()
	if err == nil {
		info.Memory = memInfo.Total / 1024 / 1024 // 转换为MB
		info.MemUsage = memInfo.UsedPercent
	}
	
	// 获取CPU使用率
	cpuPercent, err := cpu.Percent(time.Second, false)
	if err == nil && len(cpuPercent) > 0 {
		info.CPUUsage = cpuPercent[0]
	}
	
	// 获取磁盘空间
	diskInfo, err := disk.Usage("/")
	if err == nil {
		info.DiskSpace = diskInfo.Total / 1024 / 1024 / 1024 // 转换为GB
	}
	
	return info, nil
}

// GetPendingTasks 获取待处理任务
func (a *Agent) GetPendingTasks() ([]TaskInfo, error) {
	a.mu.RLock()
	if !a.registered {
		a.mu.RUnlock()
		return nil, fmt.Errorf("agent not registered")
	}
	
	workerID := a.workerID
	a.mu.RUnlock()
	
	url := fmt.Sprintf("%s/api/tasks/worker/%s/pending", strings.TrimRight(a.config.MasterURL, "/"), workerID)
	
	resp, err := a.client.Get(url)
	if err != nil {
		return nil, fmt.Errorf("failed to get pending tasks: %v", err)
	}
	defer resp.Body.Close()
	
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("failed to get pending tasks with status %d: %s", resp.StatusCode, string(body))
	}
	
	var tasks []TaskInfo
	if err := json.NewDecoder(resp.Body).Decode(&tasks); err != nil {
		return nil, fmt.Errorf("failed to decode tasks: %v", err)
	}
	
	return tasks, nil
}

// ExecuteFIOTask 执行FIO任务
func (a *Agent) ExecuteFIOTask(task TaskInfo) error {
	log.Printf("Executing FIO task %s", task.ID)
	
	// 将FIO配置转换为INI格式
	configINI, err := a.convertFIOConfigToINI(task.FIOConfig)
	if err != nil {
		log.Printf("Error converting FIO config: %v", err)
		result := TaskResult{
			Status: "failed",
			Error:  err.Error(),
		}
		return a.ReportTaskCompletion(task.ID, result)
	}
	
	// 创建临时配置文件
	configFile := fmt.Sprintf("/tmp/fio_config_%s.fio", task.ID)
	if err := a.writeFIOConfigFile(configFile, configINI); err != nil {
		log.Printf("Error writing config file: %v", err)
		result := TaskResult{
			Status: "failed",
			Error:  err.Error(),
		}
		return a.ReportTaskCompletion(task.ID, result)
	}
	
	// 执行FIO命令
	cmd := exec.Command("fio", "--output-format=json", configFile)
	
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	
	err = cmd.Run()
	
	// 清理临时配置文件
	defer func() {
		os.Remove(configFile)
	}()
	
	if err != nil {
		log.Printf("FIO execution failed: %v", err)
		result := TaskResult{
			Status: "failed",
			Error:  fmt.Sprintf("FIO execution failed: %v\nStderr: %s", err, stderr.String()),
		}
		return a.ReportTaskCompletion(task.ID, result)
	}
	
	// 解析FIO输出结果
	result := a.parseFIOOutput(stdout.String())
	result.Status = "completed"
	
	log.Printf("FIO task %s completed", task.ID)
	return a.ReportTaskCompletion(task.ID, result)
}

// convertFIOConfigToINI 将FIO配置转换为INI格式
func (a *Agent) convertFIOConfigToINI(fioConfig interface{}) (string, error) {
	configMap, ok := fioConfig.(map[string]interface{})
	if !ok {
		return "", fmt.Errorf("invalid FIO config format")
	}
	
	var iniBuilder strings.Builder
	
	// 转换全局配置参数
	for key, value := range configMap {
		if key != "jobs" { // jobs需要特别处理
			switch v := value.(type) {
			case string:
				iniBuilder.WriteString(fmt.Sprintf("%s=%s\n", key, v))
			case int:
				iniBuilder.WriteString(fmt.Sprintf("%s=%d\n", key, v))
			case float64:
				iniBuilder.WriteString(fmt.Sprintf("%s=%.0f\n", key, v))
			case bool:
				iniBuilder.WriteString(fmt.Sprintf("%s=%t\n", key, v))
			default:
				iniBuilder.WriteString(fmt.Sprintf("%s=%v\n", key, v))
			}
		}
	}
	
	return iniBuilder.String(), nil
}

// writeFIOConfigFile 写入FIO配置文件
func (a *Agent) writeFIOConfigFile(filename, content string) error {
	return os.WriteFile(filename, []byte(content), 0644)
}

// parseFIOOutput 解析FIO输出
func (a *Agent) parseFIOOutput(output string) TaskResult {
	result := TaskResult{
		Read:  make(map[string]interface{}),
		Write: make(map[string]interface{}),
	}
	
	// 简化的解析器 - 实际实现中应该使用JSON解析
	result.Read = map[string]interface{}{
		"iops":      10000,
		"bandwidth": 100000,
	}
	result.Write = map[string]interface{}{
		"iops":      5000,
		"bandwidth": 50000,
	}
	result.Latency = map[string]interface{}{
		"avg": 100,
		"max": 500,
	}
	
	return result
}

// ReportTaskCompletion 报告任务完成
func (a *Agent) ReportTaskCompletion(taskID string, result TaskResult) error {
	a.mu.RLock()
	if !a.registered {
		a.mu.RUnlock()
		return fmt.Errorf("agent not registered")
	}
	
	a.mu.RUnlock()
	
	jsonData, err := json.Marshal(result)
	if err != nil {
		return fmt.Errorf("failed to marshal result: %v", err)
	}
	
	url := fmt.Sprintf("%s/api/tasks/%s/complete", strings.TrimRight(a.config.MasterURL, "/"), taskID)
	
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create completion request: %v", err)
	}
	
	req.Header.Set("Content-Type", "application/json")
	
	resp, err := a.client.Do(req)
	if err != nil {
		return fmt.Errorf("failed to report task completion: %v", err)
	}
	defer resp.Body.Close()
	
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("failed to report task completion with status %d: %s", resp.StatusCode, string(body))
	}
	
	log.Printf("Task %s completion reported successfully", taskID)
	return nil
}

// Start 启动代理
func (a *Agent) Start() error {
	// 注册到主节点
	if err := a.RegisterWithMaster(); err != nil {
		return fmt.Errorf("failed to register with master: %v", err)
	}
	
	// 启动心跳协程
	go a.heartbeatLoop()
	
	// 启动任务处理协程
	go a.taskProcessingLoop()
	
	// 启动轮询协程
	go a.pollingLoop()
	
	log.Println("xfio agent started successfully")
	return nil
}

// heartbeatLoop 心跳循环
func (a *Agent) heartbeatLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	
	for {
		select {
		case <-ticker.C:
			if err := a.SendHeartbeat(); err != nil {
				log.Printf("Error sending heartbeat: %v", err)
			}
		}
	}
}

// pollingLoop 轮询循环
func (a *Agent) pollingLoop() {
	ticker := time.NewTicker(a.config.PollInterval)
	defer ticker.Stop()
	
	for {
		select {
		case <-ticker.C:
			// 获取待处理任务
			tasks, err := a.GetPendingTasks()
			if err != nil {
				log.Printf("Error getting pending tasks: %v", err)
				continue
			}
			
			// 将任务发送到任务通道
			for _, task := range tasks {
				select {
				case a.taskChan <- task:
					log.Printf("Added task %s to processing queue", task.ID)
				default:
					log.Printf("Task queue full, dropping task %s", task.ID)
				}
			}
		}
	}
}

// taskProcessingLoop 任务处理循环
func (a *Agent) taskProcessingLoop() {
	for task := range a.taskChan {
		log.Printf("Processing task %s", task.ID)
		if err := a.ExecuteFIOTask(task); err != nil {
			log.Printf("Error executing task %s: %v", task.ID, err)
		}
	}
}

func main() {
	var (
		masterURL    = flag.String("master", "http://localhost:8000", "Master server URL")
		token        = flag.String("token", "", "Authentication token")
		nodeName     = flag.String("name", "", "Node name")
		nodeIP       = flag.String("ip", "", "Node IP address")
		interval     = flag.Duration("interval", 30*time.Second, "Polling interval")
		insecure     = flag.Bool("insecure", false, "Allow insecure connections")
	)
	
	flag.Parse()
	
	config := &AgentConfig{
		MasterURL:    *masterURL,
		Token:        *token,
		NodeName:     *nodeName,
		NodeIP:       *nodeIP,
		PollInterval: *interval,
		Insecure:     *insecure,
	}
	
	agent := NewAgent(config)
	
	if err := agent.Start(); err != nil {
		log.Fatalf("Failed to start agent: %v", err)
	}
	
	// 等待中断信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan
	
	log.Println("Shutting down xfio agent...")
}